//
// Created by hkkzzxz on 24-8-31.
//

#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <string>
#include <vector>
#include <stdexcept>

#include "JWT.h"
#include "MySQLDatabaseModel.h"
#include "UserJWT.h"

// 对字符串进行Base64Url编码
std::string JWT::base64UrlEncode(const std::string &input) {
    std::string ret;
    int i = 0;
    int j = 0;
    unsigned char char_array_3[3];
    unsigned char char_array_4[4];

    for (auto it = input.begin(); it != input.end(); ++it) {
        char_array_3[i++] = *it;
        if (i == 3) {
            char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
            char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
            char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
            char_array_4[3] = char_array_3[2] & 0x3f;

            for (i = 0; (i < 4); i++)
                ret += base64_chars[char_array_4[i]];
            i = 0;
        }
    }

    if (i) {
        for (j = i; j < 3; j++)
            char_array_3[j] = '\0';

        char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
        char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
        char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
        char_array_4[3] = char_array_3[2] & 0x3f;

        for (j = 0; (j < i + 1); j++)
            ret += base64_chars[char_array_4[j]];

        while ((i++ < 3))
            ret += '=';
    }

    // 将Base64编码转换为Base64URL编码
    std::string base64UrlEncoded;
    for (char c: ret) {
        switch (c) {
            case '+':
                base64UrlEncoded += '-';
                break;
            case '/':
                base64UrlEncoded += '_';
                break;
            case '=': // 省略尾部的'='
                break;
            default:
                base64UrlEncoded += c;
        }
    }
    return base64UrlEncoded;
}

// Base64Url编码转化为普通字符串
std::string JWT::base64UrlDecode(const std::string &input) {
    //将Base64URL编码转换为Base64编码
    std::string base64Input = input;
    std::replace(base64Input.begin(), base64Input.end(), '-', '+');
    std::replace(base64Input.begin(), base64Input.end(), '_', '/');

    // 确保长度是4的倍数
    while (base64Input.length() % 4 != 0) {
        base64Input += '=';
    }

    // Base64解码
    BIO *bio, *b64;
    std::vector<char> buffer(input.length() + 1); // 解码后的长度不会超过输入长度

    bio = BIO_new_mem_buf(base64Input.c_str(), -1);
    b64 = BIO_new(BIO_f_base64());
    bio = BIO_push(b64, bio);

    // 不添加换行符到输出
    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);

    int decodedLength = BIO_read(bio, buffer.data(), input.length());
    buffer[decodedLength] = '\0'; // 确保字符串结束

    // 释放资源
    BIO_free_all(bio);

    return std::string(buffer.data(), decodedLength);
}

// 创建JWT签名
std::string JWT::createSignature(const std::string &headerBase64Url, const std::string &payloadBase64Url) {
    // 将头部和载荷的Base64Url编码字符串用点号连接起来，作为HMAC的输入
    std::string input = headerBase64Url + "." + payloadBase64Url;
    // 计算HMAC SHA256所需的缓冲区大小
    unsigned int len = EVP_MD_size(EVP_sha256());
    unsigned char *digest = (unsigned char *) malloc(len);
    if (digest == NULL) {
        consoleLogger->warn("Failed to allocate memory for HMAC digest.");
        return "";
    }

    // 使用HMAC和SHA256算法计算签名
    if (HMAC(EVP_sha256(), secretKey.c_str(), secretKey.size(),
             (const unsigned char *) input.c_str(), input.size(), digest, NULL) == NULL) {
        free(digest);
        consoleLogger->warn("HMAC computation failed.");
        return "";
    }

    // 将HMAC结果转换为字符串
    std::string signature = base64UrlEncode(std::string((const char *) digest, len));

    consoleLogger->info("signature:" + signature);

    // 清理HMAC分配的内存
    free(digest);

    return signature;
}

// 创建JWT
std::string JWT::create(Json::Value &payload) {
    // 设置头部
    Json::Value header;
    header["alg"] = "HS256";
    header["typ"] = "YunShi";

    // 设置负载内容
    std::ostringstream oss;
    auto duration = std::chrono::system_clock::now().time_since_epoch();
    auto seconds = std::chrono::duration_cast<std::chrono::seconds>(duration).count() + 3600;//一小时后过期
    oss << seconds;
    payload["exp"] = oss.str();
    payload["iss"] = "YunShi";

    // 对头部和载荷进行Base64编码
    std::string headerJson = header.toStyledString();
    std::string payloadJson = payload.toStyledString();
    std::string headerBase64Url = base64UrlEncode(headerJson);
    std::string payloadBase64Url = base64UrlEncode(payloadJson);

    // 签名并且进行Base64编码
    std::string signatureBase64Url = createSignature(headerBase64Url, payloadBase64Url);

    return headerBase64Url + "." + payloadBase64Url + "." + signatureBase64Url;
}

// 验证JWT
Json::Value JWT::verify(const std::string &jwt) {
    Json::Value data;

    std::vector<std::string> parts = split(jwt, '.');
    if (parts.size() != 3) {
        consoleLogger->info("Invalid JWT format");
    }

    // 获取头部、载荷和签名部分的Base64Url编码字符串
    std::string headerBase64Url = parts[0];
    std::string payloadBase64Url = parts[1];
    std::string signatureBase64Url = parts[2];

    // 重新计算签名
    std::string signature = createSignature(headerBase64Url, payloadBase64Url);

    // 比较计算出的签名和JWT中的签名是否一致
    if (signature == signatureBase64Url) {
        Json::Value payload = decodePayload(jwt);
        auto iss = payload["iss"].toStyledString();
        auto exp = payload["exp"].toStyledString();
        auto id = payload["id"].toStyledString();

        iss = remove_extra_chars(iss);
        exp = remove_extra_chars(exp);
        id = remove_extra_chars(id);

        std::time_t timestamp = std::stoll(exp);
        std::time_t now = std::time(nullptr);

        // 验证id
        std::unique_ptr<MySQLDatabaseModel> mysqlDatabaseModel = std::make_unique<MySQLDatabaseModel>();
        std::string idQuery = "SELECT * FROM user WHERE id = ?";
        std::vector<std::string> params = {id};
        bool idResult = mysqlDatabaseModel->queryBool(idQuery, params);

        // 验证负载是否符合
        if (iss == "YunShi" && timestamp > now && idResult) {
            // 查询当前用户身份
            std::string JWTQuery = "SELECT identity FROM user WHERE id = ?";
            auto userJWT = mysqlDatabaseModel->queryModel<UserJWT>(JWTQuery, params);
            if (userJWT.empty()) {
                consoleLogger->info("JWT is invalid");
                return data;
            }
            std::string identity = userJWT.front()->getIdentity();
            consoleLogger->info("JWT is valid");
            data["id"] = id;
            data["identity"] = identity;
        }
    } else {
        consoleLogger->info("JWT is invalid");
    }
    return data;
}

// 解码JWT载荷
Json::Value JWT::decodePayload(const std::string &jwt) {
    // 将JWT字符串按照'.'分割成三部分
    std::vector<std::string> parts = split(jwt, '.');
    // 如果分割后的部分不是3部分，则抛出异常
    if (parts.size() != 3) {
        consoleLogger->error("Invalid JWT format");
        fileLogger->error("Invalid JWT format");
        throw;
    }

    // 获取载荷部分的Base64Url编码字符串
    std::string payloadBase64Url = parts[1];
    // 对载荷部分进行Base64Url解码
    std::string payloadJson = base64UrlDecode(payloadBase64Url);

    Json::Value payload;
    Json::Reader reader;
    // 解析JSON字符串为Json::Value对象
    if (!reader.parse(payloadJson, payload)) {
        consoleLogger->error("Failed to parse payload");
        fileLogger->error("Failed to parse payload");
        throw;
    }
    // 返回解析后的载荷
    return payload;
}

// 将字符串按照指定的分隔符分割成字符串向量
std::vector<std::string> JWT::split(const std::string &s, char delimiter) {
    std::vector<std::string> tokens;
    std::string token;
    std::istringstream tokenStream(s);
    while (std::getline(tokenStream, token, delimiter)) {
        tokens.push_back(token);
    }
    return tokens;
}

// 规范化载荷格式
std::string JWT::remove_extra_chars(const std::string &str) {
    size_t start_pos = str.find_first_of("\"");
    if (start_pos == std::string::npos) return "";

    size_t end_pos = str.find_first_of("\"", start_pos + 1);
    if (end_pos == std::string::npos) return "";

    return str.substr(start_pos + 1, end_pos - start_pos - 1);
}
